from transformers import AutoModelForCausalLM, AutoTokenizer
import torchaudio
import pygame
import time
from openai import OpenAI 
import os
import shutil


# -- 输出划分后保存目录 --
path_save_txt = './out_story_txt/'
os.makedirs(path_save_txt, exist_ok=True)
save_file = open(path_save_txt+'0.txt', 'w', encoding='utf-8')
# -- 语音计数 --
audio_file_count = 0


client = OpenAI(
    api_key="2782135da7f0dc746cfdafde57c37276.vi2cKXc8TKbavRtb",
    base_url="https://open.bigmodel.cn/api/paas/v4/"
) 

library = '''
多人有声剧斗破苍穹完整版-p01-01-05_1_女_30,多人有声剧斗破苍穹完整版-p01-01-05_2_男_51,多人有声剧斗破苍穹完整版-p01-01-05_3_男_47,多人有声剧斗破苍穹完整版-p01-01-05_4_男_62,多人有声剧斗破苍穹完整版-p01-01-05_5_男_59,多人有声剧斗破苍穹完整版-p01-01-05_8_女_40,多人有声剧斗破苍穹完整版-p02-06-10_1_女_29,多人有声剧斗破苍穹完整版-p02-06-10_2_男_55,多人有声剧斗破苍穹完整版-p02-06-10_3_男_49,多人有声剧斗破苍穹完整版-p02-06-10_4_女_36,多人有声剧斗破苍穹完整版-p02-06-10_5_男_44,多人有声剧斗破苍穹完整版-p02-06-10_6_男_60,多人有声剧斗破苍穹完整版-p02-06-10_7_男_45,多人有声剧斗破苍穹完整版-p03-11-15_1_男_51,多人有声剧斗破苍穹完整版-p03-11-15_2_男_47,多人有声剧斗破苍穹完整版-p03-11-15_3_男_44,多人有声剧斗破苍穹完整版-p03-11-15_4_男_62,多人有声剧斗破苍穹完整版-p03-11-15_5_女_27,多人有声剧斗破苍穹完整版-p03-11-15_6_男_66,多人有声剧斗破苍穹完整版-p04-16-20_11_女_30,多人有声剧斗破苍穹完整版-p04-16-20_1_男_57,多人有声剧斗破苍穹完整版-p04-16-20_2_男_50,多人有声剧斗破苍穹完整版-p04-16-20_3_女_31,多人有声剧斗破苍穹完整版-p04-16-20_4_男_35,多人有声剧斗破苍穹完整版-p04-16-20_5_男_62,多人有声剧斗破苍穹完整版-p04-16-20_6_男_64,多人有声剧斗破苍穹完整版-p04-16-20_7_男_40,多人有声剧斗破苍穹完整版-p04-16-20_8_男_59,多人有声剧斗破苍穹完整版-p05-21-25_1_女_29,多人有声剧斗破苍穹完整版-p05-21-25_2_男_50,多人有声剧斗破苍穹完整版-p05-21-25_3_男_59,多人有声剧斗破苍穹完整版-p05-21-25_4_男_48,多人有声剧斗破苍穹完整版-p05-21-25_5_男_42,多人有声剧斗破苍穹完整版-p06-26-30_1_女_45,多人有声剧斗破苍穹完整版-p06-26-30_2_男_49,多人有声剧斗破苍穹完整版-p06-26-30_4_男_47,多人有声剧斗破苍穹完整版-p06-26-30_6_男_57, 多人有声剧斗破苍穹完整版-p06-26-30_7_女_27
'''

# prompt-3
# system = r"你是一个小说的角色匹配模型，可以根据输入的小说内容，完成以下任务：根据每个角色的性别、年龄和台词，从语料库中匹配最相关的角色，语料库为：{library}。输入格式为：[解说]：“”， [角色1_性别_年龄：]“”，[解说：]“”， [角色2_性别_年龄：]“”，等等。设定输出格式为：[解说_匹配角色]：“”， [角色1_性别_年龄_匹配角色：]“”，[解说_匹配角色：]“”， [角色2_性别_年龄_匹配角色：]“”，等等。注意，“匹配角色”只能从语料库中选择，绝不允许随意随意生成！！！"
system = f'你是一个小说的角色匹配模型，可以根据输入的小说内容，完成以下任务：根据每个角色的性别、年龄和台词，从语料库中匹配最相关的任务设定，输入的每个角色选择的匹配项前后必须一致！！带有具体名字的角色不允许从语料库中选择同一个匹配项！！！语料库为：{library}。输入格式为：[解说]：“”， [角色1_性别_年龄：]“”，[解说：]“”， [角色2_性别_年龄：]“”，等等。设定输出格式为：[解说_匹配角色]：“”， [角色1_性别_年龄_匹配角色：]“”，[解说_匹配角色：]“”， [角色2_性别_年龄_匹配角色：]“”，等等。注意，"匹配角色"尽可以从语料库中选择，绝不允许随意随意生成！！！'
input = r'[解说]：“斗之力，三段！” [解说]：望着测验魔石碑上面闪亮得甚至有些刺眼的五个大字，少年面无表情，唇角有着一抹自嘲，紧握的手掌，因为大力，而导致略微尖锐的指甲深深的刺进了掌心之中，带来一阵阵钻心的疼痛… [中年男子_男性_中年：]“萧炎，斗之力，三段！级别：低级！” [解说]：测验魔石碑之旁，一位中年男子，看了一眼碑上所显示出来的信息，语气漠然的将之公布了出来… [解说]：中年男子话刚刚脱口，便是不出意外的在人头汹涌的广场上带起了一阵嘲讽的骚动。 [未知角色_男性_青年：]“三段？嘿嘿，果然不出我所料，这个“天才”这一年又是在原地踏步！” [未知角色_男性_中年：]“哎，这废物真是把家族的脸都给丢光了。” [未知角色_男性_中年：]“要不是族长是他的父亲，这种废物，早就被驱赶出家族，任其自生自灭了，哪还有机会待在家族中白吃白喝。” [未知角色_男性_老年：]“唉，昔年那名闻乌坦城的天才少年，如今怎么落魄成这般模样了啊？” [未知角色_男性_青年：]“谁知道呢，或许做了什么亏心事，惹得神灵降怒了吧…” [解说]：周围传来的不屑嘲笑以及惋惜轻叹，落在那如木桩待在原地的少年耳中，恍如一根根利刺狠狠的扎在心脏一般，让得少年呼吸微微急促。 [解说]：少年缓缓抬起头来，露出一张有些清秀的稚嫩脸庞，漆黑的眸子木然的在周围那些嘲讽的同龄人身上扫过，少年嘴角的自嘲，似乎变得更加苦涩了。 [萧炎_男性_少年：]“这些人，都如此刻薄势力吗？或许是因为三年前他们曾经在自己面前露出过最谦卑的笑容，所以，如今想要讨还回去吧…” [解说]：苦涩的一笑，萧炎落寞的转身，安静的回到了队伍的最后一排，孤单的身影，与周围的世界，有些格格不入。 [测验人_男性_中年：]“下一个，萧媚！” [解说]：听着测验人的喊声，一名少女快速的人群中跑出，少女刚刚出场，附近的议论声便是小了许多，一双双略微火热的目光，牢牢的锁定着少女的脸颊… [解说]：少女年龄不过十四左右，虽然并算不上绝色，不过那张稚气未脱的小脸，却是蕴含着淡淡的妩媚，清纯与妩媚，矛盾的集合，让得她成功的成为了全场瞩目的焦点… [解说]：少女快步上前，小手轻车熟路的触摸着漆黑的魔石碑，然后缓缓闭上眼睛… [解说]：在少女闭眼片刻之后，漆黑的魔石碑之上再次亮起了光芒… [测验员_男性_中年：]“斗之气：七段！” [测验员_男性_中年：]“萧媚，斗之气：七段！级别:高级！” [萧媚_女性_少女：]“耶！” [解说]：听着测验员所喊出的成绩，少女脸颊扬起了得意的笑容… [未知角色_男性_青年：]“啧啧，七段斗之气，真了不起，按这进度，恐怕顶多只需要三年时间，她就能称为一名真正的斗者了吧…” [未知角色_男性_中年：]“不愧是家族中种子级别的人物啊…” [解说]：听着人群中传来的一阵阵羡慕声，少女脸颊上的笑容更是多了几分，虚荣心，这是很多女孩都无法抗拒的诱惑… [解说]：与平日里的几个姐妹互相笑谈着，萧媚的视线，忽然的透过周围的人群，停在了人群外的那一道孤单身影上… [解说]：皱眉思虑了瞬间，萧媚还是打消了过去的念头，现在的两人，已经不在同一个阶层之上，以萧炎最近几年的表现，成年后，顶多只能作为家族中的下层人员，而天赋优秀的她，则将会成为家族重点培养的强者，前途可以说是不可限量。 [萧媚_女性_少女：]“唉…” [解说]：莫名的轻叹了一口气，萧媚脑中忽然浮现出三年前那意气风发的少年，四岁练气，十岁拥有九段斗之气，十一岁突破十段斗之气，成功凝聚斗之气旋，一跃成为家族百年之内最年轻的斗者！ [解说]：当初的少年，自信而且潜力无可估量，不知让得多少少女对其春心荡漾，当然，这也包括以前的萧媚。 [解说]：然而天才的道路，貌似总是曲折的，三年之前，这名声望达到巅峰的天才少年，却是突兀的接受到了有生以来最残酷的打击，不仅辛辛苦苦修炼十数载方才凝聚的斗之气旋，一夜之间，化为乌有，而且体内的斗之气，也是随着时间的流逝，变得诡异的越来越少。 [解说]：斗之气消失的直接结果，便是导致其实力不断的后退。 [解说]：从天才的神坛，一夜跌落到了连普通人都不如的地步，这种打击，让得少年从此失魂落魄，天才之名，也是逐渐的被不屑与嘲讽所替代。 [解说]：站的越高，摔得越狠，这次的跌落，或许就再也没有爬起的机会。 [测验员_男性_中年：]“下一个，萧薰儿！” [解说]：喧闹的人群中，测试员的声音，再次响了起来。 [解说]：随着这有些清雅的名字响起，人群忽然的安静了下来，所有的视线，豁然转移。 [解说]：在众人视线汇聚之处，一位身着紫色衣裙的少女，正淡雅的站立，平静的稚嫩俏脸，并未因为众人的注目而改变分毫。 [解说]：少女清冷淡然的气质，犹如清莲初绽，小小年纪，却已初具脱俗气质，难以想象，日后若是长大，少女将会如何的倾国倾城… [解说]：这名紫裙少女，论起美貌与气质来，比先前的萧媚，无疑还要更胜上几分，也难怪在场的众人都是这般动作。 [解说]：莲步微移，名为萧薰儿的少女行到魔石碑之前，小手伸出，镶着黑金丝的紫袖滑落而下，露出一截雪白娇嫩的皓腕，然后轻触着石碑… [解说]：微微沉静，石碑之上，刺眼的光芒再次绽放。 [测验员_男性_中年：]“斗之气：九段！级别：高级！” [解说]：望着石碑之上的字体，场中陷入了一阵寂静。 [未知角色_男性_青年：]“…竟然到九段了，真是恐怖！家族中年轻一辈的第一人，恐怕非薰儿小姐莫属了。” [解说]：寂静过后，周围的少年，都是不由自主的咽了一口唾沫，眼神充满敬畏… [解说]：斗之气，每位斗者的必经之路，初阶斗之气分一至十段，当体内斗之气到达十段之时，便能凝聚斗之气旋，成为一名受人尊重的斗者！ [解说]：人群中，萧媚皱着浅眉盯着石碑前的紫裙少女，脸颊上闪过一抹嫉妒… [解说]：望着石碑上的信息，一旁的中年测验员漠然的脸庞上竟然也是罕见的露出了一丝笑意，对着少女略微恭声道： [测验员_男性_中年：]“薰儿小姐，半年之后，你应该便能凝聚斗气之旋，如果你成功的话，那么以十四岁年龄成为一名真正的斗者，你是萧家百年内的第二人！” [解说]：是的，第二人，那位第一人，便是褪去了天才光环的萧炎。 [萧薰儿_女性_少女：]“谢谢。” [解说]：少女微微点了点头，平淡的小脸并未因为他的夸奖而出现喜悦，安静的回转过身，然后在众人炽热的注目中，缓缓的行到了人群最后面的那颓废少年面前… [萧薰儿_女性_少女：]“萧炎哥哥。” [解说]：在经过少年身旁时，少女顿下了脚步，对着萧炎恭敬的弯了弯腰，美丽的俏脸上，居然露出了让周围少女为之嫉妒的清雅笑容。 [萧炎_男性_少年：]“我现在还有资格让你怎么叫么?” [解说]：望着面前这颗已经成长为家族中最璀璨的明珠，萧炎苦涩的道，她是在自己落魄后，极为少数还对自己依旧保持着尊敬的人。 [萧薰儿_女性_少女：]“萧炎哥哥，以前你曾经与薰儿说过，要能放下，才能拿起，提放自如，是自在人！” [解说]：萧薰儿微笑着柔声道，略微稚嫩的嗓音，却是暖人心肺。 [萧炎_男性_少年：]“呵呵，自在人？我也只会说而已，你看我现在的模样，象自在人吗？而且…这世界，本来就不属于我。” [解说]：萧炎自嘲的一笑，意兴阑珊的道。 [萧薰儿_女性_少女：]“萧炎哥哥，虽然并不知道你究竟是怎么回事，不过，薰儿相信，你会重新站起来，取回属于你的荣耀与尊严…” [解说]：话到此处，微顿了顿，少女白皙的俏脸，头一次露出淡淡的绯红： [萧薰儿_女性_少女：]“当年的萧炎哥哥，的确很吸引人…” [萧炎_男性_少年：]“呵呵…” [解说]：面对着少女毫不掩饰的坦率话语，少年尴尬的笑了一声，可却未再说什么，人不风流枉少年，可现在的他，实在没这资格与心情，落寞的回转过身，对着广场之外缓缓行去… [解说]：站在原地望着少年那恍如与世隔绝的孤独背影，萧薰儿踌躇了一会，然后在身后一干嫉妒的狼嚎声中，快步追了上去，与少年并肩而行…'

completion = client.chat.completions.create(
    model="glm-4-plus",  # glm-4-flash, glm-4-plus
    messages=[    
        {"role": "system", "content": system},    
        {"role": "user", "content": input} 
    ],
    top_p=0.7,
    temperature=0.9
 ) 

result_text = completion.choices[0].message.content
print(result_text)
len_text = len(result_text.split('\n'))
save_file.write(result_text)
